from pydantic import BaseModel
from typing import Optional, Literal, Dict, Any
from datetime import datetime

class StreamResponseData(BaseModel):
    """流式响应数据结构"""
    type: Literal["response_head", "reasoning", "result", "html_code", "error", "complete"]
    content: str
    message_order: Optional[int] = None
    conversation_id: Optional[str] = None
    message_id: Optional[str] = None
    frontend_message_id: Optional[str] = None
    thinking_elapsed_secs: Optional[float] = None
    accumulated_token_usage: Optional[int] = None
    status: Optional[Literal["PROCESSING", "FINISHED", "ERROR"]] = None
    
    # 扩展字段
    step: Optional[str] = None  # 当前处理步骤
    progress: Optional[float] = None  # 处理进度 0.0-1.0
    timestamp: Optional[str] = None  # 时间戳
    metadata: Optional[Dict[str, Any]] = None  # 额外元数据

class StreamResponse(BaseModel):
    """完整的流式响应格式"""
    data: StreamResponseData

# 便捷构造函数
def create_response_head(
    conversation_id: str,
    message_id: str,
    frontend_message_id: str,
    message_order: int = 1
) -> StreamResponse:
    """创建响应头"""
    return StreamResponse(
        data=StreamResponseData(
            type="response_head",
            content="开始处理请求",
            message_order=message_order,
            conversation_id=conversation_id,
            message_id=message_id,
            frontend_message_id=frontend_message_id,
            status="PROCESSING",
            timestamp=datetime.now().isoformat()
        )
    )

def create_reasoning_response(
    content: str,
    conversation_id: str,
    message_id: str,
    frontend_message_id: str,
    message_order: int,
    thinking_elapsed_secs: Optional[float] = None,
    step: Optional[str] = None
) -> StreamResponse:
    """创建推理响应"""
    return StreamResponse(
        data=StreamResponseData(
            type="reasoning",
            content=content,
            message_order=message_order,
            conversation_id=conversation_id,
            message_id=message_id,
            frontend_message_id=frontend_message_id,
            thinking_elapsed_secs=thinking_elapsed_secs,
            status="PROCESSING",
            step=step,
            timestamp=datetime.now().isoformat()
        )
    )

def create_result_response(
    content: str,
    conversation_id: str,
    message_id: str,
    frontend_message_id: str,
    message_order: int,
    accumulated_token_usage: Optional[int] = None,
    step: Optional[str] = None
) -> StreamResponse:
    """创建结果响应"""
    return StreamResponse(
        data=StreamResponseData(
            type="result",
            content=content,
            message_order=message_order,
            conversation_id=conversation_id,
            message_id=message_id,
            frontend_message_id=frontend_message_id,
            accumulated_token_usage=accumulated_token_usage,
            status="PROCESSING",
            step=step,
            timestamp=datetime.now().isoformat()
        )
    )

def create_html_code_response(
    content: str,
    conversation_id: str,
    message_id: str,
    frontend_message_id: str,
    message_order: int,
    step: Optional[str] = None
) -> StreamResponse:
    """创建HTML代码响应"""
    return StreamResponse(
        data=StreamResponseData(
            type="html_code",
            content=content,
            message_order=message_order,
            conversation_id=conversation_id,
            message_id=message_id,
            frontend_message_id=frontend_message_id,
            status="PROCESSING",
            step=step,
            timestamp=datetime.now().isoformat()
        )
    )

def create_complete_response(
    conversation_id: str,
    message_id: str,
    frontend_message_id: str,
    message_order: int,
    accumulated_token_usage: Optional[int] = None,
    thinking_elapsed_secs: Optional[float] = None
) -> StreamResponse:
    """创建完成响应"""
    return StreamResponse(
        data=StreamResponseData(
            type="complete",
            content="处理完成",
            message_order=message_order,
            conversation_id=conversation_id,
            message_id=message_id,
            frontend_message_id=frontend_message_id,
            thinking_elapsed_secs=thinking_elapsed_secs,
            accumulated_token_usage=accumulated_token_usage,
            status="FINISHED",
            timestamp=datetime.now().isoformat()
        )
    )

def create_error_response(
    error_message: str,
    conversation_id: str,
    message_id: str,
    frontend_message_id: str,
    message_order: int
) -> StreamResponse:
    """创建错误响应"""
    return StreamResponse(
        data=StreamResponseData(
            type="error",
            content=error_message,
            message_order=message_order,
            conversation_id=conversation_id,
            message_id=message_id,
            frontend_message_id=frontend_message_id,
            status="ERROR",
            timestamp=datetime.now().isoformat()
        )
    ) 